#!/usr/bin/python
# -*- coding: UTF-8 -*-
 
import os
import sys
cur_path =os.path.abspath(os.path.join(os.path.dirname('__file__'), os.path.pardir))
sys.path.append(cur_path)

from psm_conf import *

import socket
gHostname = socket.gethostname() # 获取本地主机名
gLocalIp=socket.gethostbyname(socket.getfqdn(gHostname))

import json
from kazoo.client import KazooClient
import time

import thrift

from controller_src import ps_moniter_controller
from controller_src.ttypes import *

from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from thrift.server import TServer
from kazoo.client import KazooClient

from worker.worker_src import ps_moniter_worker
from worker.worker_src.ttypes import *
from worker.worker_src.constants import *
from hulk_proxy.HulkProxyServ import Client
from scan_data.ttypes import ScanData

def log(msg):
    print(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())+"|"+msg)

def getPsmInfo(machine_type):
    # 获取线上ps_moniter worker在线信息
    info={}
    timeout = 100
    zkc = KazooClient(hosts=ZK_URL, timeout=timeout)
    zkc.start()
    if not zkc.exists(ZK_PATH+"/"+machine_type):
        log(ZK_PATH+"/"+machine_type+"不存在")
        return info
    children=zkc.get_children(ZK_PATH+"/"+machine_type)
    tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
    transport = TTransport.TFramedTransport(tsocket)
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    client = Client(protocol)
    transport.open()
    for child in children:
        tmp_hash = "psm_"+child
        ret = client.hget(tmp_hash, "ps_number")
        if not ret:
            continue # zk在线，hulk没数据，不可能的情况
        info[tmp_hash]=int(ret)
    transport.close()
    return info

class controllerHandler:
    def __init__(self):
        pass

    def runPsGroup(self, ret, input, machine_type, hulk_key, client):
        # 获取线上worker信息
        psmInfo = getPsmInfo(machine_type)
        if not psmInfo:
            ret["result"]="no %s psm_worker" % machine_type
            return
        if len(psmInfo) < RUN_ParaServ_NUMBER[machine_type]:
            ret["result"]="%s < default run ps number"%machine_type
            return
        # 加上机器类型
        jInput = json.loads(input)
        jInput["machine_type"]=machine_type

        psGroupRunRet={}
        psGroupRunRet = self.callWorker(psmInfo, json.dumps(jInput), machine_type)
        hulk_kv={}
        hulk_kv[hulk_key]=json.dumps(psGroupRunRet)# 将结果存到hulk中
        # hash固定，key记录psm_name_id_off，value以json字符串形式记录机器和psm_127.127.127.127:12345——pid|groupid|master_port
        client.hmpush(HULK_HASH, hulk_kv)
        if len(psGroupRunRet) != RUN_ParaServ_NUMBER[machine_type]+1: 
            #如果创建存在问题， 则删除
            self.stopPs(json.dumps(jInput), "offline")
            ret["result"]="len(psGroupRunRet) != offline_runPsNum+1:"

    def callWorker(self, psmInfo, input, machine_type):
        nodes={}
        i = 0
        jInput = json.loads(input)
        while i < RUN_ParaServ_NUMBER[machine_type]:
            worker_address_info = ""
            min = 9999999
            # 找出ps服务最少的机器
            for k, v in psmInfo.items():
                if v < min:
                    min = v 
                    worker_address_info = k
            # log("use "+worker_address_info)
            ip_port = worker_address_info.split('_')[1].split(':') #psm_127.127.127.127:12345
            transport = TSocket.TSocket(ip_port[0], int(ip_port[1]))
            transport = TTransport.TBufferedTransport(transport)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            client = ps_moniter_worker.Client(protocol)
            # Connect!
            transport.open()
            jInput["groupid"]=str((i+1)*1000)
            ret = client.runPs(json.dumps(jInput)) # {0:pid|master_port}
            log("Create ParaServ: "+worker_address_info+" "+str(ret))
            if 0 not in ret:
                break
            psmInfo.pop(worker_address_info) #正常返回则删除这台机器, 不再创建psServer
            #返回ps_moniter机器上的ParaServ的 pid|groupid|master_port
            nodes[worker_address_info] = ret[0]
            if i+1 == RUN_ParaServ_NUMBER[machine_type]:
                # psServer启动完成，启动psProxy
                try:
                    ret = client.runPsProxy(json.dumps(jInput))
                    log("Create PSProxy: "+worker_address_info+" "+str(ret))
                    if 0 in ret:
                        #返回ps_moniter机器上的PSProxy的pid
                        nodes[worker_address_info+":"+str("proxy")] = ret[0] 
                except:
                    log("runPsProxy error!")
            transport.close()
            i+=1
        return nodes

    def runOnlinePs(self, input):
        log("runOnlinePs"+"|in|"+input)
        ret={}
        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        # 检查model_id，必须在线
        jInput=json.loads(input)
        model_name_id=jInput["model_name"]+"_"+str(jInput["model_id"])
        ret['model_id']=str(jInput["model_id"])

        hulk_key = "psm_"+model_name_id
        checkModelId=client.hget(HULK_HASH, hulk_key+"_realtime")#realtime 必须在线
        
        if not checkModelId:
            log("toBeOnline but hulk_key:%s doesn`t exit: "%(hulk_key+"_realtime"))
            ret["result"]=checkModelId
            return json.dumps(ret)              
          
        self.runPsGroup(ret, input, "online", hulk_key+"_online", client)
        if "result" in ret:
            transport.close()
            return json.dumps(ret)

        #关闭offline集群
        self.stopPs(input, "realtime")

        ret["zk_url_online"]=ZK_URL
        ret["zk_path_online"]="{NEW_ZKPATH}/psProxy/{MODEL_NAME_ID}_{MACHINE_TYPE}".format(NEW_ZKPATH=ZK_PATH, MODEL_NAME_ID=model_name_id, MACHINE_TYPE="online")

        log("runOnlinePs"+"|out|"+input)
        return json.dumps(ret)
        
    def runRealtimePs(self, input):
        log("runRealtimePs"+"|in|"+input)
        ret={}
        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        jInput=json.loads(input)
        model_name_id=jInput["model_name"]+"_"+str(jInput["model_id"])
        ret['model_id']=str(jInput["model_id"])

        hulk_key = "psm_"+model_name_id
        checkModelId=client.hget(HULK_HASH, hulk_key+"_offline")#offline 必须在线
        
        if not checkModelId:
            log("toBeOnline but hulk_key:%s doesn`t exit: "%(hulk_key+"_offline"))
            ret["result"]=checkModelId
            return json.dumps(ret)              
          
        self.runPsGroup(ret, input, "realtime", hulk_key+"_realtime", client)
        if "result" in ret:
            transport.close()
            return json.dumps(ret)   

        #关闭offline集群
        self.stopPs(input, "offline")

        ret["result"]="OK"
        ret["zk_url_realtime"]=ZK_URL
        ret["zk_path_realtime"]="{NEW_ZKPATH}/psProxy/{MODEL_NAME_ID}_{MACHINE_TYPE}".format(NEW_ZKPATH=ZK_PATH, MODEL_NAME_ID=model_name_id, MACHINE_TYPE="realtime")

        log("runRealtimePs"+"|out|"+input)
        return json.dumps(ret)

    def runOfflinePs(self, input):
        log("runOfflinePs"+"|in|"+input)
        ret={}
        jInput=json.loads(input)
        model_name_id=jInput["model_name"]+"_"+str(jInput["model_id"])
        ret['model_id']=str(jInput["model_id"])
        ret['model_name']=str(jInput["model_name"])
        hulk_key = "psm"+'_'+model_name_id+"_offline"

        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        # 检查model_id
        checkModelId=client.hget(HULK_HASH, hulk_key)

        if checkModelId:
            log("runOfflinePs but hulk_key:%s exits: "%hulk_key)
            ret["result"]="RUNNING"
            ret["zk_url"]=ZK_URL
            ret["zk_path"]=ZK_PATH+"/psProxy/"+model_name_id+"_offline"
            log("runOfflinePs"+"|out|"+input+"|"+json.dumps(ret))
            return json.dumps(ret)

        self.runPsGroup(ret, input, "offline", hulk_key, client)
        transport.close()

        if  "result" not in ret:
            ret["result"]="OK"
            ret["zk_url"]=ZK_URL
            ret["zk_path"]=ZK_PATH+"/psProxy/"+model_name_id+"_offline"

        log("runOfflinePs"+"|out|"+input+"|"+json.dumps(ret))
        return json.dumps(ret)

    def stopPs(self, input, machine_type):
        log("stopPs|in|"+input+"|machine_type="+machine_type)
        ret={}
        jInput=json.loads(input)
        model_name_id=jInput["model_name"]+"_"+str(jInput["model_id"])
        ret['model_id']=str(jInput["model_id"])
        hulk_machine_type=machine_type
        hulk_key = "psm"+'_'+model_name_id+"_"+hulk_machine_type
        # 查询hulk，确认该model_name_id没有在线
        tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
        transport = TTransport.TFramedTransport(tsocket)
        protocol = TBinaryProtocol.TBinaryProtocol(transport)
        client = Client(protocol)
        transport.open()
        psGroupInfo=client.hget(HULK_HASH, hulk_key)
        if not psGroupInfo:
            ret["result"]=model_name_id+" no runing ps process"
            log("stopPs|out|"+input+"|"+json.dumps(ret))
            return json.dumps(ret)
        json_psGroupInfo=json.loads(psGroupInfo)
        log("These process will be stopped: "+psGroupInfo)
        # 关闭进程    
        for k,v in json_psGroupInfo.items():
            #k => psm_127.127.127.127:12345、slave_127.127.127.127:12345
            #v => pid|groupid|master_port
            ip_port = k.split('_')[1].split(':') 
            ttransport = TSocket.TSocket(ip_port[0], int(ip_port[1]))
            ttransport = TTransport.TBufferedTransport(ttransport)
            pprotocol = TBinaryProtocol.TBinaryProtocol(ttransport)
            cclient = ps_moniter_worker.Client(pprotocol)
            # Connect!
            ttransport.open()
            pid=v.split("|")[0]
            #master_port=v.split("|")[1]
            #groupid=v.split("|")[2]
            cclient.stopPs(str(pid)+"|"+model_name_id+"|"+("proxy" if k.find("proxy")!=-1 else "ParaServ"))
            ttransport.close()
        # 清空hulk
        client.hdel(HULK_HASH, hulk_key)
        transport.close()
        # 清空zk
        zkc = KazooClient(hosts=ZK_URL, timeout=100)
        zkc.start()
        # 删除ParaServ注册地址
        machine_type = "offline" if "machine_type" not in jInput else jInput["machine_type"]
        ParaServ_path=ZK_ParaServ_REGISTE_PATH+model_name_id+"_"+machine_type
        if not zkc.exists(ParaServ_path):
            log("No: "+ParaServ_path)
        else: 
            log("Delete ParaServ registe path: "+ParaServ_path)
            zkc.delete(ParaServ_path, recursive=True)
        # 删除tcmcp_psProxy地址
        tcmcp_psProxy=ZK_PATH+"/psProxy/"+model_name_id+"_"+machine_type
        if not zkc.exists(tcmcp_psProxy):
            log("No: "+tcmcp_psProxy)
        else: 
            log("Delete tcmcp_psProxy path: "+tcmcp_psProxy)
            zkc.delete(tcmcp_psProxy, recursive=True)
        
        ret["result"]="OK"
        log("stopPs|out|"+input+"|machine_type="+machine_type)
        return json.dumps(ret)

    def getInfo(self, input):
        log("getInfo|in|"+input)
        jInput=json.loads(input)
        ret={}
        if jInput["cmd"] == "list":
            # 查询每台worker机器上ParaServ的数量
            machine_types=["offline", "realtime", "online"]
            for machine_type in machine_types:
                psmInfo = getPsmInfo(machine_type)
                log("%s:%d"%(machine_type, len(psmInfo)))
                if not psmInfo:
                    ret[machine_type]="0"
                    continue
                for key in psmInfo.keys():
                    ret[machine_type+"_"+key.split("_")[1].split(":")[0]]=str(psmInfo[key])
        elif jInput["cmd"] == "info":
            workers=[]
            if "ip" in jInput:
                workers.append("input_"+jInput["ip"])
            else:
                machine_types=["offline", "realtime", "online"]
                for machine_type in machine_types:
                    psmInfo = getPsmInfo(machine_type)
                    if psmInfo:
                        for key in psmInfo.keys():
                            workers.append(machine_type+"_"+key.split("_")[1].split(":")[0]) # psm_127.0.0.1:1234
            for worker in workers:
                ip=worker.split('_')[1]
                ttransport = TSocket.TSocket(ip, PORT_PSM_WORKER)
                ttransport = TTransport.TBufferedTransport(ttransport)
                pprotocol = TBinaryProtocol.TBinaryProtocol(ttransport)
                cclient = ps_moniter_worker.Client(pprotocol)
                ttransport.open()
                ret[worker]=cclient.getInfo("a")
                ttransport.close()
        elif jInput["cmd"] == "model":
            tsocket = TSocket.TSocket(HULK_HOST, HULK_PORT)
            transport = TTransport.TFramedTransport(tsocket)
            protocol = TBinaryProtocol.TBinaryProtocol(transport)
            client = Client(protocol)
            transport.open()
            hulk_key=""
            if "model_name" in jInput and "model_id" in jInput and "machine_type" in jInput:
                hulk_key="psm_"+jInput["model_name"]+"_"+jInput["model_id"]+"_"+jInput["machine_type"]
            if hulk_key:
                psGroupInfo=client.hget(HULK_HASH, hulk_key)
                ret[hulk_key]=psGroupInfo
            else:
                hashKey_psGroupInfo=client.hgetall(HULK_HASH)
                for k,v in hashKey_psGroupInfo.items():
                    ret[k]=v
        else:
            ret["result"]="getInfo input error"
        log("getInfo|out|"+str(ret))
        return  ret


handler = controllerHandler()
processor = ps_moniter_controller.Processor(handler)
transport = TSocket.TServerSocket(gLocalIp, PORT_PSM_CONTROLLER)
tfactory = TTransport.TBufferedTransportFactory()
pfactory = TBinaryProtocol.TBinaryProtocolFactory()
server = TServer.TThreadPoolServer(processor, transport, tfactory, pfactory)
#server = TServer.TSimpleServer(processor, transport, tfactory, pfactory)
log("Starting ps_moniter controller server...")
server.serve()

